package cn.tank.data.sync.kafka;

import cn.tank.data.sync.core.MessageAcceptService;
import cn.tank.data.sync.core.common.Message;
import cn.tank.data.sync.core.common.Subscriber;
import cn.tank.data.sync.kafka.common.KafkaManager;
import org.apache.kafka.clients.consumer.Consumer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Collections;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicReference;

/************************************************
 * @Description:kafka消息消费实现
 * @Author: tancyong
 * @Version: 1.0
 ************************************************/
public class KafkaMessageConsumer implements MessageAcceptService {

    private static final Logger loggr = LoggerFactory.getLogger(KafkaMessageConsumer.class);

    private static final String DATA_TYPE_HEADER = "DATATYPE";

    private final KafkaManager kafkaManager;

    private Map<String, ConsumerThread> consumerThreadMap = new ConcurrentHashMap();

    public KafkaMessageConsumer(KafkaManager kafkaManager) {

        this.kafkaManager = kafkaManager;

    }


    public void createSubscriber(String topic, String groupId, Subscriber subscriber) {
        try {
            Consumer<String, String> consumer = kafkaManager.getConsumer(groupId);
            consumer.subscribe(Collections.singleton(topic));

            loggr.info("订阅者{},成功监听主题{}", subscriber.getName(), topic);
            ConsumerThread consumerThread = new ConsumerThread(consumer, topic, subscriber);
            consumerThreadMap.put(subscriber.getId(), consumerThread);
            consumerThread.start();

        } catch (Exception e) {
            e.printStackTrace();

            ConsumerThread consumerThread = consumerThreadMap.get(subscriber.getId());
            if (consumerThread != null) {
                consumerThread.shutdown();
                consumerThreadMap.remove(subscriber.getId());
            }
        }
    }


    public void receiveMessage(Subscriber subscriber) {
        ConsumerThread consumerThread = consumerThreadMap.get(subscriber.getId());

        if (consumerThread != null) {
            consumerThread.shutdown();
            consumerThreadMap.remove(subscriber.getId());
        }
    }


    class ConsumerThread extends Thread {

        private final String topic;
        private final Subscriber subscriber;
        private final Consumer<String, String> consumer;

        private boolean isRunning = true;

        public ConsumerThread(Consumer<String, String> consumer, String topic, Subscriber subscriber) {

            this.topic = topic;
            this.subscriber = subscriber;
            this.consumer = consumer;
        }

        @Override
        public void run() {

            while (isRunning) {
                consumer.poll(1000).forEach(record -> {
                    AtomicReference<String> dataType = new AtomicReference<>("");
                    record.headers().forEach(header -> {
                        if (header.key().equalsIgnoreCase(DATA_TYPE_HEADER)) {
                            dataType.set(new String(header.value()));
                        }
                    });

                    Message message = new Message(dataType.get(), new String(record.value()));
                    receiveMessage(topic, subscriber, message);
                });
            }

            consumer.close();
            consumer.wakeup();
        }

        public void shutdown() {
            this.isRunning = false;
        }
    }

}

